import os
import requests
from pathlib import Path
from dotenv import load_dotenv

load_dotenv() # 加载.env文件中的环境变量


def generate_description_from_image_zhipu(prompt, image64, zhipu_api_key):
    """
    用智谱 glm-4v 实现“图片→描述”
    :param prompt: 用户提示
    :param image64: base64 编码的图片
    :param zhipu_api_key: 智谱控制台生成的 API Key
    :return: 描述文本
    """
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {zhipu_api_key}"
    }

    payload = {
        "model": "glm-4v",          # 带视觉能力的模型
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image64}"}}
                ]
            }
        ],
        "max_tokens": 300
    }

    resp = requests.post(
        "https://open.bigmodel.cn/api/paas/v4/chat/completions",
        headers=headers,
        json=payload,
        timeout=30
    )
    resp.raise_for_status()
    return resp.json()['choices'][0]['message']['content']

def generate_image_zhipuai(prompt: str, zhipuai_api_key: str, save_path: str = "generated_image.png") -> str:
    """
    使用智谱AI生成图片（最终修正版）
    :param prompt: 图片描述文本
    :param zhipuai_api_key: 有效的API Key
    :param save_path: 图片保存路径
    :return: 图片本地路径或None（失败时）
    """
    # 确保保存目录存在
    Path(save_path).parent.mkdir(parents=True, exist_ok=True)

    # API请求配置
    url = "https://open.bigmodel.cn/api/paas/v4/images/generations"
    headers = {
        "Authorization": f"Bearer {zhipuai_api_key}",
        "Content-Type": "application/json"
    }
    payload = {
        "model": "cogview-3",
        "prompt": prompt,
        "n": 1,
        "size": "1024x1024",
        "response_format": "url"
    }

    try:
        # 1. 创建图片生成任务
        response = requests.post(url, headers=headers, json=payload, timeout=30)
        response.raise_for_status()
        response_data = response.json()

        # 2. 检查返回数据结构
        if not response_data.get("data") or not isinstance(response_data["data"], list):
            raise ValueError("API返回数据格式异常")

        # 3. 获取图片URL
        image_url = response_data["data"][0]["url"]
        print(f"获取到图片URL: {image_url}")

        # 4. 下载图片
        image_response = requests.get(image_url, timeout=30)
        image_response.raise_for_status()

        # 5. 保存图片
        with open(save_path, "wb") as f:
            f.write(image_response.content)

        print(f"图片已成功保存到: {save_path}")
        return save_path

    except Exception as e:
        print(f"发生错误: {str(e)}")
        print("完整API响应:", response_data if 'response_data' in locals() else "无响应数据")
        return None

# 使用示例
if __name__ == "__main__":
    # 替换为你的真实API Key
    API_KEY = os.getenv("ZHIPUAI_API_KEY")

    result = generate_image_zhipuai(
        prompt="一只穿着宇航服的柴犬在月球表面漫步，科幻风格，4K高清",
        zhipuai_api_key=API_KEY,
        save_path="./data/astronaut_dog.png"
    )

    if result:
        print(f"图片生成成功！保存路径: {result}")
    else:
        print("图片生成失败，请检查错误信息")

# 示例调用
# if __name__ == "__main__":
#     api_key = os.getenv("ZHIPUAI_API_KEY")
#     with open("D:/tmp/images.png", "rb") as f:
#         img_b64 = base64.b64encode(f.read()).decode()
#     desc = generate_description_from_image_zhipu("请描述这张图片", img_b64, api_key)
#     print("智谱生成描述:", desc)